# coding = utf-8

'''
将数据转换到台湾成功大学模型所需要的样本
'''
import os
import pydicom
import numpy as np
from pathlib2 import Path
import matplotlib.pyplot as plt


def normalize(vol):
    hu_max = 250
    hu_min = -200
    vol = np.clip(vol, hu_min, hu_max)

    mxval = np.max(vol)
    mnval = np.min(vol)
    volume_norm = (vol - mnval) / max(mxval - mnval, 1e-3)

    return volume_norm

def get_pixels_hu(scans):
    image = np.stack([s.pixel_array for s in scans])
    # Convert to int16 (from sometimes int16),
    # should be possible as values should always be low enough (<32k)
    image = image.astype(np.int16)

    # Set outside-of-scan pixels to 1
    # The intercept is usually -1024, so air is approximately 0
    image[image == -2000] = 0

    # Convert to Hounsfield units (HU)
    intercept = scans[0].RescaleIntercept
    slope = scans[0].RescaleSlope

    if slope != 1:
        image = slope * image.astype(np.float64)
        image = image.astype(np.int16)

    image += np.int16(intercept)
    return np.array(image, dtype=np.int16)


def readImage(path):
    index_2_image = {}
    for item in os.listdir(path):
        file_name = os.path.join(path, item)
        dcm = pydicom.dcmread(file_name)
        hu = get_pixels_hu([dcm])
        hu = normalize(hu.squeeze())
        index_2_image[item.split("_")[1].zfill(3)] = hu
    return index_2_image


def readMask(path):
    index_2_mask = {}
    liver_mask_path = os.path.join(path, "liver")
    tumor_mask_list = []
    for name in os.listdir(path):
        if "livertumor" in name:
            tumor_mask_list.append(os.path.join(path,name))
    for file_index in os.listdir(liver_mask_path):
        liver_mask_file = os.path.join(liver_mask_path, file_index)
        liver_mask = pydicom.dcmread(liver_mask_file)
        mask = liver_mask.pixel_array
        mask_division = np.max(mask)
        if mask_division == 0:
            mask_division = 255
        mask = mask/mask_division
        print(mask_division)
        mask = mask.astype(np.uint8)
        for tumor_mask_path in tumor_mask_list:
            tumor_mask_file = os.path.join(tumor_mask_path, file_index)
            assert os.path.exists(tumor_mask_file), tumor_mask_file+" not exist"
            tumor_mask = pydicom.dcmread(tumor_mask_file)
            tumor_mask = tumor_mask.pixel_array
            tumor_division = np.max(tumor_mask)
            if tumor_division == 0:
                tumor_division = 255
            tumor_mask = tumor_mask / tumor_division
            print(tumor_division)
            tumor_mask = tumor_mask.astype(np.uint8)
            tumor_size = (tumor_mask == 1).sum()
            liver_size = (mask == 1).sum()
            inter = (mask[tumor_mask == 1] == 1).sum()
            #assert (tumor_size-inter) == 0, "{},liver size:{}, tumor size:{}, inter:{}, diff:{}".format(tumor_mask_file,liver_size, tumor_size, inter,
            #                                                                                         tumor_size-inter)
            mask[tumor_mask == 1] = 2
        index_2_mask[file_index.split("_")[1].zfill(3)] = mask
    return index_2_mask


def read_data():
    root = "/datasets/3Dircadb/3Dircadb1"
    destination = "/datasets/3Dircadb/chengkun_only_liver"
    for patient_id in os.listdir(root):
        patient_path = os.path.join(root, patient_id)
        image_path = os.path.join(patient_path, "PATIENT_DICOM")
        index_2_image = readImage(image_path)
        mask_path = os.path.join(patient_path, "MASKS_DICOM")
        index_2_mask = readMask(mask_path)
        case_id = str(int(patient_id.split(".")[1]) - 1).zfill(5)
        destination_patient_path = os.path.join(destination, "case_{}".format(case_id))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path/"imaging"
        destination_mask_path = destination_patient_path/"segmentation"
        print(destination_image_path,destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        z = 0
        for key in sorted(index_2_image.keys()):
            assert key in index_2_mask.keys(), "{} not in mask".format(key)
            image = index_2_image[key]
            mask = index_2_mask[key]
            assert image.shape == mask.shape, "image shape != mask_shape, patient_id:{}, key:{}, image.shape:{}, mask.shape:{},".\
                format(patient_id, key, image.shape, mask.shape)
            #if np.max(mask) == 0:
            #    continue
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(z).zfill(3))), image)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(z).zfill(3))), mask)
            z += 1

        print("finish {}".format(patient_id))

#对转换之后的数据做一些简单的校验
def chenkung_data_validate():
    root = "/datasets/3Dircadb/3Dircadb1"
    destination = "/datasets/3Dircadb/chengkung"

    for patient_id in os.listdir(root):
        patient_path = os.path.join(root, patient_id)
        image_path = os.path.join(patient_path, "PATIENT_DICOM")
        mask_path = os.path.join(patient_path, "MASKS_DICOM")
        case_id = str(int(patient_id.split(".")[1]) - 1).zfill(5)
        destination_patient_path = os.path.join(destination, "case_{}".format(case_id))
        destination_image_path = os.path.join(destination_patient_path, "imaging")
        destination_mask_path = os.path.join(destination_patient_path, "segmentation")

        slice_number = len(os.listdir(image_path))
        dst_image_number = len(os.listdir(destination_image_path))
        dst_mask_number = len(os.listdir(destination_mask_path))

        assert slice_number == dst_image_number and dst_image_number == dst_mask_number, "slice_number != dst_image_number != dst_mask_number, " \
                                                                                         "{} != {} != {}".format(slice_number, dst_image_number, dst_mask_number)



        tumor_mask_list = []
        for name in os.listdir(mask_path):
            if "livertumor" in name:
                tumor_mask_list.append(os.path.join(mask_path, name))

        mask_list = None
        for item in os.listdir(destination_mask_path):
            mask_file = os.path.join(destination_mask_path, item)
            mask = np.load(mask_file)
            mask = mask.reshape((1,mask.shape[0], mask.shape[1]))
            if mask_list is None:
                mask_list = mask
            else:
                mask_list = np.concatenate([mask_list,mask], axis=0)

        unique_mask = np.unique(mask_list)
        #print()

        assert 0 in unique_mask, patient_id
        assert 1 in unique_mask, patient_id
        if len(tumor_mask_list) > 0:
            assert 2 in unique_mask, patient_id

        print(patient_id, case_id, slice_number, unique_mask)


def origion_mask_validate():
    root = "/datasets/3Dircadb/3Dircadb1"
    patient_id = "3Dircadb1.2"
    patient_path = os.path.join(root,patient_id)
    mask_path = os.path.join(patient_path, "MASKS_DICOM")
    liver_mask_path = os.path.join(mask_path, "liver")
    for file_name in os.listdir(liver_mask_path):
        liver_mask_file = os.path.join(liver_mask_path,file_name)
        liver_mask = pydicom.dcmread(liver_mask_file)
        print(np.min(liver_mask.pixel_array), np.max(liver_mask.pixel_array))

#对于转换之后的数据进行相应的展示
def chenkung_data_visul():
    root = "/datasets/3Dircadb/3Dircadb1"
    patient_id = "3Dircadb1.1"
    destination = "/datasets/3Dircadb/chengkung"
    patient_path = os.path.join(root, patient_id)
    mask_path = os.path.join(patient_path, "MASKS_DICOM")
    liver_mask_path = os.path.join(mask_path, "liver")
    case_id = "case_{}".format(str(int(patient_id.split(".")[1]) - 1).zfill(5))
    dst_patient_path = os.path.join(destination, case_id)

    tumor_mask_list = []
    for name in os.listdir(mask_path):
        if "livertumor" in name:
            tumor_mask_list.append(os.path.join(mask_path, name))

    for file_index in os.listdir(liver_mask_path):
        liver_mask_file = os.path.join(liver_mask_path, file_index)
        liver_mask = pydicom.dcmread(liver_mask_file)
        if np.max(liver_mask.pixel_array) <= 0:
            continue
        dst_file_index = file_index.split("_")[1].zfill(3)

        dst_mask_file = os.path.join(dst_patient_path, "segmentation//{}.npy".format(dst_file_index))
        mask = np.load(dst_mask_file)
        dst_image_file = os.path.join(dst_patient_path, "imaging//{}.npy".format(dst_file_index))
        image = np.load(dst_image_file)

        tumor_mask = np.zeros(liver_mask.pixel_array.shape)

        for tumor_mask_path in tumor_mask_list:
            tumor_mask_file = os.path.join(tumor_mask_path, file_index)
            assert os.path.exists(tumor_mask_file), tumor_mask_file+" not exist"
            tumor_mask_dcm = pydicom.dcmread(tumor_mask_file)
            tumor_mask_dcm = tumor_mask_dcm.pixel_array
            tumor_mask[tumor_mask_dcm>0] = 1

        print(file_index, case_id, dst_file_index)

        plt.subplot(2, 2, 1)
        plt.imshow(liver_mask.pixel_array, cmap="gray")
        plt.subplot(2,2,2)
        plt.imshow(tumor_mask, cmap="gray")
        plt.subplot(2, 2, 3)
        plt.imshow(image, cmap="gray")
        plt.subplot(2, 2, 4)
        plt.imshow(mask)
        plt.show()









if __name__ == '__main__':
  read_data()